#!/usr/bin/env python3
"""
AI应用环境变量管理器
用于管理.ai/applications目录下程序的环境变量和系统路径
"""

import os
import sys
import json
import subprocess
import shutil
from pathlib import Path
from typing import List, Dict, Optional

class EnvironmentManager:
    def __init__(self, apps_dir: Optional[Path] = None):
        if apps_dir is None:
            # 默认为.ai/applications目录
            current_dir = Path.cwd()
            self.apps_dir = current_dir / ".ai" / "applications"
        else:
            self.apps_dir = Path(apps_dir)

        self.config_file = self.apps_dir / "env-config.json"
        self.shell_configs = {
            "bash": "~/.bashrc",
            "zsh": "~/.zshrc",
            "fish": "~/.config/fish/config.fish",
            "windows": "环境变量设置"
        }

    def load_config(self) -> Dict:
        """加载环境配置"""
        if not self.config_file.exists():
            return {"applications": {}, "installed_at": None}

        try:
            with open(self.config_file, 'r', encoding='utf-8') as f:
                return json.load(f)
        except:
            return {"applications": {}, "installed_at": None}

    def save_config(self, config: Dict):
        """保存环境配置"""
        self.apps_dir.mkdir(parents=True, exist_ok=True)

        config["updated_at"] = str(Path.cwd())

        with open(self.config_file, 'w', encoding='utf-8') as f:
            json.dump(config, f, indent=2, ensure_ascii=False)

    def discover_applications(self) -> List[Path]:
        """发现applications目录下的可执行程序"""
        if not self.apps_dir.exists():
            return []

        applications = []

        # 查找Python脚本
        for py_file in self.apps_dir.glob("*.py"):
            if py_file.name != "env-manager.py":  # 排除管理器本身
                applications.append(py_file)

        # 查找可执行文件
        if os.name != 'nt':  # 非Windows系统
            for exe_file in self.apps_dir.glob("*"):
                if exe_file.is_file() and os.access(exe_file, os.X_OK):
                    applications.append(exe_file)
        else:  # Windows系统
            for exe_file in self.apps_dir.glob("*.exe"):
                applications.append(exe_file)

        return applications

    def register_application(self, app_path: Path, command_name: str = None) -> bool:
        """注册应用程序"""
        if not app_path.exists():
            print(f"❌ 应用程序不存在: {app_path}")
            return False

        if command_name is None:
            command_name = app_path.stem

        config = self.load_config()

        app_info = {
            "name": command_name,
            "path": str(app_path),
            "type": "python" if app_path.suffix == ".py" else "executable",
            "installed_at": str(Path.cwd())
        }

        config["applications"][command_name] = app_info
        self.save_config(config)

        print(f"✅ 已注册应用程序: {command_name}")
        return True

    def unregister_application(self, command_name: str) -> bool:
        """注销应用程序"""
        config = self.load_config()

        if command_name in config["applications"]:
            del config["applications"][command_name]
            self.save_config(config)
            print(f"✅ 已注销应用程序: {command_name}")
            return True
        else:
            print(f"❌ 应用程序不存在: {command_name}")
            return False

    def setup_environment(self) -> bool:
        """设置环境变量"""
        config = self.load_config()

        if not config["applications"]:
            print("❌ 没有已注册的应用程序")
            return False

        # 创建包装脚本目录
        wrapper_dir = Path.home() / ".ai-apps"
        wrapper_dir.mkdir(exist_ok=True)

        print(f"📁 创建包装脚本目录: {wrapper_dir}")

        success_count = 0

        for app_name, app_info in config["applications"].items():
            try:
                if app_info["type"] == "python":
                    success = self._create_python_wrapper(app_name, app_info, wrapper_dir)
                else:
                    success = self._create_executable_wrapper(app_name, app_info, wrapper_dir)

                if success:
                    success_count += 1

            except Exception as e:
                print(f"❌ 创建 {app_name} 包装脚本失败: {e}")

        # 更新PATH环境变量
        if success_count > 0:
            self._update_path_env(wrapper_dir)
            print(f"✅ 成功设置 {success_count} 个应用程序")

            # 显示使用说明
            self._show_usage_instructions(wrapper_dir)

            return True
        else:
            print("❌ 没有成功设置任何应用程序")
            return False

    def _create_python_wrapper(self, app_name: str, app_info: Dict, wrapper_dir: Path) -> bool:
        """创建Python应用的包装脚本"""
        app_path = Path(app_info["path"])

        if os.name == 'nt':  # Windows
            wrapper_file = wrapper_dir / f"{app_name}.bat"
            content = f"""@echo off
python "{app_path}" %*
"""
        else:  # Linux/Mac
            wrapper_file = wrapper_dir / app_name
            content = f"""#!/bin/bash
python "{app_path}" "$@"
"""

        with open(wrapper_file, 'w', encoding='utf-8') as f:
            f.write(content)

        if os.name != 'nt':
            os.chmod(wrapper_file, 0o755)

        print(f"✅ 创建包装脚本: {wrapper_file}")
        return True

    def _create_executable_wrapper(self, app_name: str, app_info: Dict, wrapper_dir: Path) -> bool:
        """创建可执行文件的包装脚本"""
        app_path = Path(app_info["path"])

        if os.name == 'nt':  # Windows
            wrapper_file = wrapper_dir / f"{app_name}.bat"
            content = f"""@echo off
"{app_path}" %*
"""
        else:  # Linux/Mac
            wrapper_file = wrapper_dir / app_name
            content = f"""#!/bin/bash
"{app_path}" "$@"
"""

        with open(wrapper_file, 'w', encoding='utf-8') as f:
            f.write(content)

        if os.name != 'nt':
            os.chmod(wrapper_file, 0o755)

        print(f"✅ 创建包装脚本: {wrapper_file}")
        return True

    def _update_path_env(self, wrapper_dir: Path):
        """更新PATH环境变量"""
        wrapper_path = str(wrapper_dir)

        if os.name == 'nt':  # Windows
            try:
                import winreg

                # 获取当前PATH
                key = winreg.OpenKey(winreg.HKEY_CURRENT_USER, "Environment", 0, winreg.KEY_ALL_ACCESS)
                current_path, _ = winreg.QueryValueEx(key, "PATH")
                winreg.CloseKey(key)

                # 检查是否已在PATH中
                if wrapper_path not in current_path:
                    new_path = f"{current_path};{wrapper_path}"

                    # 更新注册表
                    key = winreg.OpenKey(winreg.HKEY_CURRENT_USER, "Environment", 0, winreg.KEY_ALL_ACCESS)
                    winreg.SetValueEx(key, "PATH", 0, winreg.REG_EXPAND_SZ, new_path)
                    winreg.CloseKey(key)

                    print(f"✅ 已添加到Windows环境变量PATH: {wrapper_path}")
                    print("⚠️  请重启命令提示符以使更改生效")
                else:
                    print(f"✅ 目录已在PATH中: {wrapper_path}")

            except Exception as e:
                print(f"❌ 更新Windows环境变量失败: {e}")
                print(f"请手动将 {wrapper_path} 添加到PATH环境变量")
        else:  # Linux/Mac
            shell_rc_files = [
                Path.home() / ".bashrc",
                Path.home() / ".zshrc",
                Path.home() / ".profile"
            ]

            added_to_file = False
            for rc_file in shell_rc_files:
                if rc_file.exists():
                    try:
                        content = rc_file.read_text()
                        export_line = f'export PATH="$PATH:{wrapper_path}"'

                        if export_line not in content:
                            with open(rc_file, 'a', encoding='utf-8') as f:
                                f.write(f"\n# AI Applications\n{export_line}\n")
                            print(f"✅ 已添加到 {rc_file.name}")
                            added_to_file = True
                            break

                    except Exception as e:
                        print(f"❌ 更新 {rc_file.name} 失败: {e}")

            if not added_to_file:
                print(f"⚠️  请手动将以下内容添加到您的shell配置文件:")
                print(f'export PATH="$PATH:{wrapper_path}"')
            else:
                print("⚠️  请重启终端或执行 source ~/.bashrc 以使更改生效")

    def _show_usage_instructions(self, wrapper_dir: Path):
        """显示使用说明"""
        print("\n📖 使用说明:")
        print("=" * 50)
        print(f"包装脚本目录: {wrapper_dir}")
        print("现在您可以在命令行中直接使用以下命令:")

        config = self.load_config()
        for app_name in config["applications"]:
            print(f"  • {app_name}")

        print("\n如果命令无法使用，请:")
        print("1. 重启命令提示符/终端")
        print("2. 或者手动刷新环境变量")
        if os.name == 'nt':
            print("   Windows: 重启命令提示符")
        else:
            print("   Linux/Mac: source ~/.bashrc 或重新打开终端")

    def list_applications(self):
        """列出已注册的应用程序"""
        config = self.load_config()

        if not config["applications"]:
            print("❌ 没有已注册的应用程序")
            return

        print("📋 已注册的应用程序:")
        print("=" * 50)

        for app_name, app_info in config["applications"].items():
            print(f"🔹 {app_name}")
            print(f"   路径: {app_info['path']}")
            print(f"   类型: {app_info['type']}")
            print(f"   安装位置: {app_info['installed_at']}")
            print()

    def auto_discover_and_register(self) -> int:
        """自动发现并注册应用程序"""
        applications = self.discover_applications()

        if not applications:
            print("❌ 在applications目录中没有发现应用程序")
            return 0

        registered_count = 0

        for app_path in applications:
            command_name = app_path.stem
            if self.register_application(app_path, command_name):
                registered_count += 1

        print(f"\n🎉 自动发现并注册了 {registered_count} 个应用程序")
        return registered_count

    def sync_from_local(self, original_cwd: str) -> bool:
        """从本地同步applications到系统环境"""
        print("🔄 同步本地applications到系统环境...")

        # 发现本地applications目录中的程序
        local_apps_dir = Path(original_cwd) / ".ai" / "applications"

        if not local_apps_dir.exists():
            print("❌ 本地applications目录不存在")
            return False

        print(f"📁 扫描本地目录: {local_apps_dir}")

        # 发现Python脚本
        py_files = list(local_apps_dir.glob("*.py"))
        exe_files = []

        if os.name != 'nt':
            exe_files = [f for f in local_apps_dir.glob("*")
                        if f.is_file() and os.access(f, os.X_OK)]
        else:
            exe_files = list(local_apps_dir.glob("*.exe"))

        total_apps = len(py_files) + len(exe_files)

        if total_apps == 0:
            print("❌ 本地目录中没有发现可执行程序")
            return False

        print(f"🔍 发现 {total_apps} 个应用程序:")

        # 注册所有发现的应用程序
        registered_count = 0

        for py_file in py_files:
            if py_file.name != "env-manager.py":
                if self.register_application(py_file):
                    registered_count += 1

        for exe_file in exe_files:
            if self.register_application(exe_file):
                registered_count += 1

        if registered_count > 0:
            print(f"\n✅ 已注册 {registered_count} 个应用程序")
            print("🔧 正在设置环境变量...")
            return self.setup_environment()
        else:
            print("❌ 没有成功注册任何应用程序")
            return False

def main():
    """主函数"""
    import argparse

    parser = argparse.ArgumentParser(
        description="AI应用环境变量管理器",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例:
  %(prog)s auto-discover              # 自动发现并注册应用程序
  %(prog)s register app.py myapp      # 注册应用程序
  %(prog)s setup                     # 设置环境变量
  %(prog)s list                      # 列出已注册的应用程序
  %(prog)s unregister myapp          # 注销应用程序
  %(prog)s sync                      # 从本地同步到系统环境
        """
    )

    subparsers = parser.add_subparsers(dest='command', help='可用命令')

    # auto-discover命令
    subparsers.add_parser('auto-discover', help='自动发现并注册应用程序')

    # register命令
    register_parser = subparsers.add_parser('register', help='注册应用程序')
    register_parser.add_argument('app_path', help='应用程序路径')
    register_parser.add_argument('command_name', nargs='?', help='命令名称（可选）')

    # setup命令
    subparsers.add_parser('setup', help='设置环境变量')

    # list命令
    subparsers.add_parser('list', help='列出已注册的应用程序')

    # unregister命令
    unregister_parser = subparsers.add_parser('unregister', help='注销应用程序')
    unregister_parser.add_argument('command_name', help='命令名称')

    # sync命令
    sync_parser = subparsers.add_parser('sync', help='从本地同步到系统环境')
    sync_parser.add_argument('--cwd', help='本地工作目录路径')

    args = parser.parse_args()

    if not args.command:
        parser.print_help()
        return

    try:
        env_manager = EnvironmentManager()

        if args.command == 'auto-discover':
            env_manager.auto_discover_and_register()

        elif args.command == 'register':
            app_path = Path(args.app_path)
            env_manager.register_application(app_path, args.command_name)

        elif args.command == 'setup':
            env_manager.setup_environment()

        elif args.command == 'list':
            env_manager.list_applications()

        elif args.command == 'unregister':
            env_manager.unregister_application(args.command_name)

        elif args.command == 'sync':
            cwd = args.cwd or os.getcwd()
            env_manager.sync_from_local(cwd)

    except Exception as e:
        print(f"❌ 执行失败: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()